import scipy.io
import numpy as np
import math
import random

## 算法实现

def Average(train_set):
	result = 0.0
	cnt = 0
	for arr in train_set:
		cnt += 1
		result += arr[2]
	return result / cnt


def InerProduct(v1, v2):
	return v1@v2.T


def PredictScore(av, bu, bi, pu, qi):
	pScore = av + bu + bi + InerProduct(pu, qi)
	if pScore < 1:
		pScore = 1
	elif pScore > 5:
		pScore = 5
	return pScore

def Validate(probe_set, av, bu, bi, pu, qi):
	cnt = 0
	rmse = 0.0		
	for arr in probe_set:
		cnt += 1
		uid = arr[0] - 1
		iid = arr[1] - 1
		pScore = PredictScore(av, bu[uid], bi[iid], pu[uid], qi[iid])
		tScore = arr[2]
		rmse += (tScore - pScore) * (tScore - pScore)
	# print(cnt)
	return math.sqrt(rmse / cnt)

def MF(train_set, probe_set, factorNum, learnRate, regularization, maxiter):
    print(factorNum, learnRate, regularization)
    bu = [0.0 for i in range(num_p)] #list,初始化列表,bu表示第u个用户的偏离程度
    bu = np.array(bu)
    bi = [0.0 for i in range(num_m)] #list,初始化列表,bi表示第i部电影的偏离程度
    bi = np.array(bi)
    temp = math.sqrt(factorNum) #temp是浮点型数字，根号10
    qi = [[(0.1 * random.random() / temp) for j in range(factorNum)] for i in range(num_m)]	
    pu = [[(0.1 * random.random() / temp)  for j in range(factorNum)] for i in range(num_p)]
    qi = np.array(qi)
    pu = np.array(pu)
    print("initialization end \n start training \n")#初始化完毕
    # train model
    preRmse = 1000000.0
    for step in range(maxiter):#迭代100次
        for arr in train_set:#根据下面的代码来看,训练集包括三个元素用户id,电影id和对应的评分作为训练集
            uid = int(arr[0]) - 1
            iid = int(arr[1]) - 1
            score = int(arr[2])	
            #print("test_RMSE in step %d %d\n" %(uid, iid))		
            prediction = PredictScore(averageScore, bu[uid], bi[iid], pu[uid], qi[iid])#进行预测，这是一个迭代的过程
            #print("test_RMSE in step1\n")				
            eui = score - prediction
            #更新用户和商品列表参数
            bu[uid] += learnRate * (eui - regularization * bu[uid])
            bi[iid] += learnRate * (eui - regularization * bi[iid])	
            temp=pu[uid]
            pu[uid]+=learnRate*(eui*qi[iid]-regularization*pu[uid])
            qi[iid]+=learnRate*(eui*temp-regularization*qi[iid])
        ## 学习率逐渐下降
        learnRate *= 0.99
        curRmse = Validate(probe_set, averageScore, bu, bi, pu, qi)
        print("test_RMSE in step %d: %f" %(step, curRmse))
        preRmse = curRmse


if __name__ == '__main__':
    ## 加载数据集
    ### 1、movielens-1M Datasets
    num_p = 6040
    num_m = 3952
    mat = scipy.io.loadmat("./datasets/MovieLens 1M Dataset.mat")
    # ### 2、movieslens-100k Datasets
    # num_p = 6040
    # num_m = 3952
    # mat = scipy.io.loadmat("./datasets/MovieLens 100K Dataset.mat")

    # ### 3、MovieLens Latest Datasets
    # num_p = 6040
    # num_m = 3952
    # mat = scipy.io.loadmat("./datasets/ml-latest-small.mat")
    probe_set = mat['probe_vec']
    train_set = mat['train_vec']
    print(type(probe_set))
    print(type(train_set))
    averageScore = Average(train_set)
    print(averageScore)
    MF(train_set, probe_set, factorNum = 40, learnRate = 0.01, regularization = 0.05, maxiter = 200)